/*
Copyright 2025 PerfXLab (Beijing) Technologies Co., Ltd.

Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at

    http://www.apache.org/licenses/LICENSE-2.0

Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
*/
#include "test_spmv_vnec.hpp"
#define ntest 2000
#define nwarm 1

inline void test_spmv_vnec_fp32(const std::string &input_path, const std::string &result_path, bool iscsc)
{
    Builder<OpenSpB_index, OpenSpB_index> builder(input_path);
    CSRGraph<OpenSpB_index, OpenSpB_index> g = builder.MakeGraph();

    OpenSpB_index nrows, ncols, nvals;

    nrows = static_cast<OpenSpB_index>(g.num_nodes());
    ncols = static_cast<OpenSpB_index>(g.num_nodes());
    nvals = static_cast<OpenSpB_index>(g.num_edges());

    OpenSpB_index *row_indices;
    OpenSpB_index *col_indices;
    float *values = new float[nvals];

    row_indices = new OpenSpB_index[nvals];
    col_indices = g.out_destValues();
    OpenSpB_index *offsets = g.out_offsets();
    OpenSpB_index index = 0;
    for (OpenSpB_index i = 0; i < nrows; i++)
    {

        for (OpenSpB_index j = 0; j < offsets[i + 1] - offsets[i]; j++)
        {
            if (index >= nvals)
            {
                std::cerr << "Error index >= nvals" << index << " " << nvals << std::endl;
                exit(-1);
            }
            values[index] = i;
            row_indices[index++] = static_cast<OpenSpB_index>(i);
        }
    }

    OpenSpB_init();
    OpenSpB_object_format_ext format;
    format.matrix_format = OpenSpB_csr;
    OpenSpB_Global_Option_set_ext(OpenSpB_matrix_format, format);
    OpenSpB_matrix G = NULL;
    OpenSpB_info info;
    SAFE_CALL(OpenSpB_matrix_new(&G, OpenSpB_FP32, nrows, ncols));
    OpenSpB_matrix_build_FP32(G, row_indices, col_indices, values, nvals);
    OpenSpB_vector x;
    SAFE_CALL(OpenSpB_vector_new(&x, OpenSpB_FP32, ncols));
    for (OpenSpB_index j = 0; j < ncols; j++)
    {
        OpenSpB_vector_set_element_FP32(x, j, j);
    }
    OpenSpB_vector result_benchmark;
    SAFE_CALL(OpenSpB_vector_new(&result_benchmark, OpenSpB_FP32, nrows));
    OpenSpB_vector result_opt;
    SAFE_CALL(OpenSpB_vector_new(&result_opt, OpenSpB_FP32, nrows));
    struct timeval opt_start, opt_end;
    for (OpenSpB_index i = 0; i < 1; i++)
    {
        OpenSpB_mxv_FP32_serial(result_benchmark, G, x);
    }
    //---------------------------------------------------------------------------------------------------------------------------
    if (nrows < 3000000)
    {
        OpenSpB_VNEC_type vnec_t = IRD_VNEC(G, OpenSpB_FP32);
        if (vnec_t == OpenSpB_VNEC_D)
        {
            printf("-D is selected\n");
            VNEC_D_FP32 *mat_thd = Build_VNEC_D_FP32(G);
            // warm up
            for (OpenSpB_index i = 0; i < nwarm; i++)
            {
                SpMV_VNEC_D_FP32(result_opt, G, x, mat_thd);
            }
            gettimeofday(&opt_start, NULL);
            for (OpenSpB_index i = 0; i < ntest; i++)
            {
                SpMV_VNEC_D_FP32(result_opt, G, x, mat_thd);
            }
            gettimeofday(&opt_end, NULL);
            VNEC_D_free_fp32(mat_thd);
        }
        else
        {
            printf("-S is selected\n");
            VNEC_S_FP32 *mat_thd = Build_VNEC_S_FP32(G);
            // warm up
            for (OpenSpB_index i = 0; i < nwarm; i++)
            {
                SpMV_VNEC_S_FP32(result_opt, G, x, mat_thd);
            }
            gettimeofday(&opt_start, NULL);
            for (OpenSpB_index i = 0; i < ntest; i++)
            {
                SpMV_VNEC_S_FP32(result_opt, G, x, mat_thd);
            }
            gettimeofday(&opt_end, NULL);
            VNEC_S_free_fp32(mat_thd);
        }
    }
    else
    {
        printf("-L is selected\n");
        VNEC_L *mat_thd = Build_VNEC_L(G);
        // warm up
        for (OpenSpB_index i = 0; i < nwarm; i++)
        {
            SpMV_VNEC_L_FP32(result_opt, G, x, mat_thd);
        }
        gettimeofday(&opt_start, NULL);
        for (OpenSpB_index i = 0; i < ntest; i++)
        {
            SpMV_VNEC_L_FP32(result_opt, G, x, mat_thd);
        }
        gettimeofday(&opt_end, NULL);
        free(mat_thd->diagonal_start);
        free(mat_thd->diagonal_end);
        free(mat_thd->nz_indices);
        free(mat_thd->thread_coord_start);
        free(mat_thd->thread_coord_end);
        free(mat_thd);
    }
    long opt_timeuse = 1000000 * (opt_end.tv_sec - opt_start.tv_sec) + opt_end.tv_usec - opt_start.tv_usec;
    printf("\n ------------------------------------------------------------------------------");
    printf("\n SpMV overall time elapsed: %f ms\n", opt_timeuse / 1000.0 / ntest);
    printf(" ------------------------------------------------------------------------------ \n");
    float diff = 0.0;
    for (OpenSpB_index i = 0; i < nrows; ++i)
    {
        float value_benchmark, value_opt;
        OpenSpB_vector_extract_element_FP32(&value_benchmark, result_benchmark, i);
        OpenSpB_vector_extract_element_FP32(&value_opt, result_opt, i);
        diff = fabs(value_benchmark - value_opt);
        ASSERT_LE(diff, value_benchmark * 1e-4);
    }

    OpenSpB_matrix_free(&G);
    OpenSpB_vector_free(&x);
    OpenSpB_vector_free(&result_opt);
    OpenSpB_vector_free(&result_benchmark);
}

inline void test_spmv_vnec_fp64(const std::string &input_path, const std::string &result_path, bool iscsc)
{
    Builder<OpenSpB_index, OpenSpB_index> builder(input_path);
    CSRGraph<OpenSpB_index, OpenSpB_index> g = builder.MakeGraph();

    OpenSpB_index nrows, ncols, nvals;

    nrows = static_cast<OpenSpB_index>(g.num_nodes());
    ncols = static_cast<OpenSpB_index>(g.num_nodes());
    nvals = static_cast<OpenSpB_index>(g.num_edges());

    OpenSpB_index *row_indices;
    OpenSpB_index *col_indices;
    double *values = new double[nvals];

    row_indices = new OpenSpB_index[nvals];
    col_indices = g.out_destValues();
    OpenSpB_index *offsets = g.out_offsets();
    OpenSpB_index index = 0;
    for (OpenSpB_index i = 0; i < nrows; i++)
    {

        for (OpenSpB_index j = 0; j < offsets[i + 1] - offsets[i]; j++)
        {
            if (index >= nvals)
            {
                std::cerr << "Error index >= nvals" << index << " " << nvals << std::endl;
                exit(-1);
            }
            values[index] = i;
            row_indices[index++] = static_cast<OpenSpB_index>(i);
        }
    }
    OpenSpB_init();
    OpenSpB_object_format_ext format;
    format.matrix_format = OpenSpB_csr;
    OpenSpB_Global_Option_set_ext(OpenSpB_matrix_format, format);
    OpenSpB_matrix G = NULL;
    OpenSpB_info info;
    SAFE_CALL(OpenSpB_matrix_new(&G, OpenSpB_FP64, nrows, ncols));
    OpenSpB_matrix_build_FP64(G, row_indices, col_indices, values, nvals);
    OpenSpB_vector x;
    SAFE_CALL(OpenSpB_vector_new(&x, OpenSpB_FP64, ncols));
    for (OpenSpB_index j = 0; j < ncols; j++)
    {
        OpenSpB_vector_set_element_FP64(x, j, j);
    }
    OpenSpB_vector result_benchmark;
    SAFE_CALL(OpenSpB_vector_new(&result_benchmark, OpenSpB_FP64, nrows));
    OpenSpB_vector result_opt;
    SAFE_CALL(OpenSpB_vector_new(&result_opt, OpenSpB_FP64, nrows));
    struct timeval opt_start, opt_end;
    for (OpenSpB_index i = 0; i < 1; i++)
    {
        OpenSpB_mxv_FP64_serial(result_benchmark, G, x);
    }
    //---------------------------------------------------------------------------------------------------------------------------
    if (nrows < 3000000)
    {
        OpenSpB_VNEC_type vnec_t = IRD_VNEC(G, OpenSpB_FP64);
        if (vnec_t == OpenSpB_VNEC_D)
        {
            printf("-D is selected\n");
            VNEC_D_FP64 *mat_thd = Build_VNEC_D_FP64(G);
            // warm up
            for (OpenSpB_index i = 0; i < nwarm; i++)
            {
                SpMV_VNEC_D_FP64(result_opt, G, x, mat_thd);
            }
            gettimeofday(&opt_start, NULL);
            for (OpenSpB_index i = 0; i < ntest; i++)
            {
                SpMV_VNEC_D_FP64(result_opt, G, x, mat_thd);
            }
            gettimeofday(&opt_end, NULL);
            VNEC_D_free_fp64(mat_thd);
        }
        else
        {
            printf("-S is selected\n");
            VNEC_S_FP64 *mat_thd = Build_VNEC_S_FP64(G);
            // warm up
            for (OpenSpB_index i = 0; i < nwarm; i++)
            {
                SpMV_VNEC_S_FP64(result_opt, G, x, mat_thd);
            }
            gettimeofday(&opt_start, NULL);
            for (OpenSpB_index i = 0; i < ntest; i++)
            {
                SpMV_VNEC_S_FP64(result_opt, G, x, mat_thd);
            }
            gettimeofday(&opt_end, NULL);
            VNEC_S_free_fp64(mat_thd);
        }
    }
    else
    {
        printf("-L is selected\n");
        VNEC_L *mat_thd = Build_VNEC_L(G);
        // warm up
        for (OpenSpB_index i = 0; i < nwarm; i++)
        {
            SpMV_VNEC_L_FP64(result_opt, G, x, mat_thd);
        }
        gettimeofday(&opt_start, NULL);
        for (OpenSpB_index i = 0; i < ntest; i++)
        {
            SpMV_VNEC_L_FP64(result_opt, G, x, mat_thd);
        }
        gettimeofday(&opt_end, NULL);
        free(mat_thd->diagonal_start);
        free(mat_thd->diagonal_end);
        free(mat_thd->nz_indices);
        free(mat_thd->thread_coord_start);
        free(mat_thd->thread_coord_end);
        free(mat_thd);
    }
    long opt_timeuse = 1000000 * (opt_end.tv_sec - opt_start.tv_sec) + opt_end.tv_usec - opt_start.tv_usec;
    printf("\n ------------------------------------------------------------------------------");
    printf("\n VNEC_FP64 mxv overall time elapsed: %f ms\n", opt_timeuse / 1000.0 / ntest);
    printf(" ------------------------------------------------------------------------------ \n");
    double diff = 0.0;
    for (OpenSpB_index i = 0; i < nrows; ++i)
    {
        double value_benchmark, value_opt;
        OpenSpB_vector_extract_element_FP64(&value_benchmark, result_benchmark, i);
        OpenSpB_vector_extract_element_FP64(&value_opt, result_opt, i);
        diff = fabs(value_benchmark - value_opt);
        ASSERT_LE(diff, value_benchmark * 1e-8);
    }

    OpenSpB_matrix_free(&G);
    OpenSpB_vector_free(&x);
    OpenSpB_vector_free(&result_opt);
    OpenSpB_vector_free(&result_benchmark);
}

inline void test_read_fp32(const std::string &sg_path, const std::string &m_path)
{
    Builder<OpenSpB_index, OpenSpB_index> builder(sg_path);
    CSRGraph<OpenSpB_index, OpenSpB_index> g = builder.MakeGraph();
    OpenSpB_index nrows, ncols, nvals;
    nrows = static_cast<OpenSpB_index>(g.num_nodes());
    ncols = static_cast<OpenSpB_index>(g.num_nodes());
    nvals = static_cast<OpenSpB_index>(g.num_edges());

    OpenSpB_index *row_indices;
    OpenSpB_index *col_indices;
    float *values = new float[nvals];

    row_indices = new OpenSpB_index[nvals];
    col_indices = g.out_destValues();
    OpenSpB_index *offsets = g.out_offsets();
    OpenSpB_index index = 0;
    for (OpenSpB_index i = 0; i < nrows; i++)
    {

        for (OpenSpB_index j = 0; j < offsets[i + 1] - offsets[i]; j++)
        {
            if (index >= nvals)
            {
                std::cerr << "Error index >= nvals" << index << " " << nvals << std::endl;
                exit(-1);
            }
            values[index] = 1;
            row_indices[index++] = static_cast<OpenSpB_index>(i);
        }
    }

    OpenSpB_init();
    OpenSpB_object_format_ext format;
    format.matrix_format = OpenSpB_csr;
    OpenSpB_Global_Option_set_ext(OpenSpB_matrix_format, format);
    OpenSpB_matrix G = NULL;
    OpenSpB_info info;
    SAFE_CALL(OpenSpB_matrix_new(&G, OpenSpB_FP32, nrows, ncols));
    OpenSpB_matrix_build_FP32(G, row_indices, col_indices, values, nvals);

    std::vector<OpenSpB_index> row_indices_;
    std::vector<OpenSpB_index> col_indices_;
    std::vector<float> values_;
    OpenSpB_index nrows_, ncols_, nvals_;
    const char *mtx_path = m_path.c_str();

    readMtx_coo<float>(mtx_path, &row_indices_, &col_indices_, &values_, &nrows_, &ncols_, &nvals_, 1, false);

    OpenSpB_matrix G_ = NULL;
    SAFE_CALL(OpenSpB_matrix_new(&G_, OpenSpB_FP32, nrows_, ncols_));
    OpenSpB_matrix_build_FP32(G_, &row_indices_[0], &col_indices_[0], (float *)&values_[0], nvals_);

    OpenSpB_vector x;
    SAFE_CALL(OpenSpB_vector_new(&x, OpenSpB_FP32, ncols));
    for (OpenSpB_index j = 0; j < ncols; j++)
    {
        OpenSpB_vector_set_element_FP32(x, 1, j);
    }
    OpenSpB_vector result_mm;
    SAFE_CALL(OpenSpB_vector_new(&result_mm, OpenSpB_FP32, nrows));
    OpenSpB_vector result_sg;
    SAFE_CALL(OpenSpB_vector_new(&result_sg, OpenSpB_FP32, nrows));

    OpenSpB_mxv_FP32_serial(result_sg, G, x);
    OpenSpB_mxv_FP32_serial(result_mm, G_, x);
    // float diff = 0.0;
    float sum_mm = 0.0;
    float sum_sg = 0.0;
    for (OpenSpB_index i = 1; i < nrows; ++i)
    {
        float value_sg, value_mm;
        OpenSpB_vector_extract_element_FP32(&value_sg, result_sg, i);
        OpenSpB_vector_extract_element_FP32(&value_mm, result_mm, i);
        // diff = fabs(value_sg - value_mm);
        // ASSERT_LE(diff, value_mm * 1e-4);
        sum_mm += value_mm;
        sum_sg += value_sg;
        if(value_mm != value_sg)
        {
            printf("\n (%ld)  mm=%f sg = %f  \n",i,value_mm,value_sg);
        }
        ASSERT_EQ(value_mm, value_sg);
    }
    // printf("\n mm=%f sg = %f  \n",sum_mm,sum_sg);

    OpenSpB_matrix_free(&G);
    OpenSpB_vector_free(&x);
    OpenSpB_vector_free(&result_sg);
    OpenSpB_vector_free(&result_mm);

    // float diff = 0.0;
    // float sum_mm = 0.0;
    // float sum_sg = 0.0;
    // for (OpenSpB_index i = 0; i < 1; ++i)
    // {
    //     for (OpenSpB_index j = 0; j < 1; ++j)
    //     {
    //         float value_mm, value_sg;
    //         OpenSpB_matrix_extract_element_FP32(&value_mm, G_, i, j);
    //         OpenSpB_matrix_extract_element_FP32(&value_sg, G, i, j);
    //         sum_mm += value_mm;
    //         sum_sg += value_sg;
    //         // diff = fabs(value_mm - value_sg);
    //         // ASSERT_LE(diff, value_mm * 1e-4);
    //         // if(value_mm != value_sg)
    //         // {
    //         //     printf("\n (%ld,%ld)  mm=%f sg = %f  \n",i,j,value_mm,value_sg);
    //         // }
    //     }
    // }
    // printf("\n mm=%f sg = %f  \n",sum_mm,sum_sg);

    // OpenSpB_matrix_free(&G);
    OpenSpB_matrix_free(&G_);

}
#define BATCH_TEST(case_, name)	\
TEST(mxv_test_read_FP32, case_)	\
{	\
    std::string input_file1(MATRIX_SG_FULL_PATH(name));	\
    std::string input_file2(MATRIX_FULL_PATH(name));	\
    test_read_fp32(input_file1, input_file2);		\
}	\
TEST(mxv_test_VNEC_FP32, case_)	\
{	\
    std::string input_file(MATRIX_SG_FULL_PATH(name));	\
    std::string output_file(RESULT_FULL_PATH(name));	\
    test_spmv_vnec_fp32(input_file, output_file, false);\
}	\
TEST(mxv_test_VNEC_FP64, case_)	\
{	\
    std::string input_file(MATRIX_SG_FULL_PATH(name));	\
    std::string output_file(RESULT_FULL_PATH(name));	\
    test_spmv_vnec_fp64(input_file, output_file, false);\
}	
BATCH_TEST(amazon0302, amazon0302)
BATCH_TEST(ca_CondMat, ca-CondMat)
BATCH_TEST(email_EuAll, email-EuAll)
//BATCH_TEST(Oregon1, Oregon-1)
//BATCH_TEST(Oregon2, Oregon-2)
//BATCH_TEST(p2pGnutella04, p2p-Gnutella04)
//BATCH_TEST(p2pGnutella05, p2p-Gnutella05)
//BATCH_TEST(p2pGnutella06, p2p-Gnutella06)
//BATCH_TEST(p2pGnutella08, p2p-Gnutella08)
//BATCH_TEST(p2pGnutella09, p2p-Gnutella09)
//BATCH_TEST(p2pGnutella24, p2p-Gnutella24)
//BATCH_TEST(p2pGnutella25, p2p-Gnutella25)
//BATCH_TEST(p2pGnutella30, p2p-Gnutella30)
//BATCH_TEST(p2pGnutella31, p2p-Gnutella31)
BATCH_TEST(roadNetPA, roadNet-PA)
BATCH_TEST(roadNetTX, roadNet-TX)
BATCH_TEST(socEpinions1, soc-Epinions1)
//BATCH_TEST(socLiveJournal1, soc-LiveJournal1) //slow
BATCH_TEST(socsignepinions, soc-sign-epinions)
BATCH_TEST(socsignSlashdot081106, soc-sign-Slashdot081106)
BATCH_TEST(socsignSlashdot090216, soc-sign-Slashdot090216)
BATCH_TEST(socsignSlashdot090221, soc-sign-Slashdot090221)
BATCH_TEST(socSlashdot0811, soc-Slashdot0811)
//BATCH_TEST(socSlashdot0902, soc-Slashdot0902) //segement failed
BATCH_TEST(webNotreDame, web-NotreDame)
//BATCH_TEST(wikiTalk, wiki-Talk)
//BATCH_TEST(wikiVote, wiki-Vote)
